import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms

from config import Config
from model import UNet
from utils import denormalize, WaterMark

best_loss = 1000000
img_size = 512

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])


def get_image(image):
    output = (denormalize(image.permute((0, 2, 3, 1)).
                          detach().
                          to('cpu').
                          numpy()) * 255).astype('uint8')
    output = output[0]

    output = Image.fromarray(output).convert('RGB')
    return output


class Nosie2Noisetrain(object):
    def __init__(self):
        self.args = Config()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)
            torch.cuda.manual_seed_all(self.args.seed)
        else:
            torch.manual_seed(self.args.seed)

        self.device = torch.device('cuda' if use_cuda else 'cpu')

        '''
        定义选择模型
        '''
        self.model = UNet().to(self.device)

        '''
        根据需要加载与训练的模型权重参数
        '''
        if self.args.resume and self.args.pretrained_model:
            model_dict = self.model.state_dict()
            checkpoint = torch.load(self.args.pretrained_model)
            pretrained_dict = checkpoint['model_state_dict']
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if
                               np.shape(model_dict[k]) == np.shape(v)}
            model_dict.update(pretrained_dict)
            self.model.load_state_dict(model_dict, strict=True)
            print("Restoring the weight from pretrained-weight file \nFinished to load the weight")

        self.test()
        torch.cuda.empty_cache()
        print("finish model training")

    def test(self, x=1):
        self.model.eval()
        with torch.no_grad():
            data = Image.open("./000017.jpg")
            data = data.resize((img_size, img_size))
            temp1 = data
            data = WaterMark()(data)

            temp2 = data
            data = test_transform(data)

            data = data.unsqueeze(0)
            output = data.to(self.device)

            for i in range(x):
                output = self.model(output)
                output = nn.Tanh()(output)
                output = get_image(output)
                output = test_transform(output).to(self.device)
                output=output.unsqueeze(0)

            result = get_image(output)

            blank = Image.new(mode='RGB', size=(3 * img_size, img_size), color=(255, 255, 255))

            blank.paste(temp1, (0, 0))
            blank.paste(temp2, (img_size, 0))
            blank.paste(result, (img_size * 2, 0))
            blank.show()
            blank.save("result.jpg")


Nosie2Noisetrain()
